import torch
from .PoissonEncoding import PoissonEncoding

def spike_time_to_seq(stime: torch.Tensor, tsteps: int = 1000) -> torch.Tensor:
    spike_seq = torch.zeros(tsteps+1)
    for i in stime:
        spike_seq[int(i * 10)] = 1
    return spike_seq

def spike_seq_to_time(seq: torch.Tensor, tsteps: int = 1000) -> torch.Tensor:
    stime = torch.nonzero(seq).squeeze() * 0.1
    if stime.dim() == 0:
        stime = stime.unsqueeze(0)
    return stime



def gen_seq_learning_task(num_input: int, rate_input: int = 200,
                          num_output: int = 1, rate_output: int = 100,
                          tsteps: int = 1000):
    """
    生成脉冲序列学习任务

    参数:
    num_input(int): 输入脉冲的条数
    rate_input(int): 输入脉冲的频率,频率越高，生成脉冲越密集
    num_output(int): 输出脉冲的条数
    rate_output(int): 输出脉冲频率
    tsteps(int): 模拟的时间长度，默认为1000

    返回值:
    torch.Tensor: 输入脉冲序列, shape(tsteps, num_input)
    torch.Tensor: 输出脉冲序列, shape(tsteps, num_output)
    """
    input_spike_list = []
    
    for i in range(num_input):
        pei = PoissonEncoding(rate_input)
        stime = pei.spike_encoding(tsteps * 0.1)     #这里泊松编码器默认时间步长0.1
        input_spike_list.append(spike_time_to_seq(stime, tsteps))
    input_spike_tensor = torch.stack(input_spike_list).T

    output_spike_list = []
    for i in range(num_output):
        peo = PoissonEncoding(rate_output)
        stime = peo.spike_encoding(tsteps * 0.1)
        output_spike_list.append(spike_time_to_seq(stime, tsteps))
    output_spike_tensor = torch.stack(output_spike_list).T
    return input_spike_tensor, output_spike_tensor